package com.chukun.flink.dataset.partition;

import com.chukun.flink.dataset.bean.Trade;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;

import java.util.ArrayList;
import java.util.List;

/**
 * @author chukun
 * @version 1.0.0
 * @description 数据物理分区的基本操作
 * @createTime 2022年05月29日 14:42:00
 */
public class PartitionTemplate {

    public static void main(String[] args) throws Exception {

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(3);

        final String flag = " 分区策略前子任务名称:";

        List<Trade> list = new ArrayList<Trade>();
        list.add(new Trade("185001", 1199, "2018"));
        list.add(new Trade("155002", 1111, "2019"));
        list.add(new Trade("155002", 20, "2019"));
        list.add(new Trade("185001", 2899, "2018"));
        list.add(new Trade("138003", 19, "2019"));
        list.add(new Trade("138003", 399, "2020"));
        list.add(new Trade("138003", 39, "2020"));
        list.add(new Trade("138003", 99, "2020"));

        DataSet<Trade> dataSource = env.fromCollection(list);

        DataSet<Trade> mapResult = dataSource.map(new RichMapFunction<Trade, Trade>() {
            @Override
            public Trade map(Trade value) {
                String taskNameWithSubtasks = getRuntimeContext().getTaskNameWithSubtasks();
                int indexOfThisSubtask = getRuntimeContext().getIndexOfThisSubtask();
                System.out.println("元素值:" + value + flag + taskNameWithSubtasks + " ,子任务编号:" + indexOfThisSubtask);
                return value;
            }
        }).setParallelism(3);

//        mapResult = mapResult.rebalance();
        mapResult = mapResult.partitionByHash("cardNum");
//        mapResult = mapResult.partitionByRange("cardNum");
//        mapResult = mapResult.sortPartition("cardNum", Order.ASCENDING);

        mapResult = mapResult.map(new RichMapFunction<Trade, Trade>() {
            @Override
            public Trade map(Trade value) {
                System.out.println("元素值:" + value + " 分区策略后子任务名称:" + getRuntimeContext().getTaskNameWithSubtasks()
                        + " ,子任务编号:" + getRuntimeContext().getIndexOfThisSubtask());
                return value;
            }
        });
        mapResult.print("输出结果");

        env.execute("PartitionTemplate");
    }
}
